package com.sccba.uop.demo.util.algorithm;

import org.bouncycastle.jce.provider.BouncyCastleProvider;
import org.bouncycastle.pqc.math.linearalgebra.ByteUtils;

import javax.crypto.*;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;
import java.security.*;

/**
 * SM4算法工具
 */
public class SM4Utils {

    private static final String SM4 = "SM4";
    private static final String MODE = "CBC";
    private static final String PADDING = "PKCS7Padding";

    private static final boolean NEED_IV = true;
    private static final String IV_HEX_STRING = "00000000000000000000000000000000";

    static {
        if (Security.getProvider(BouncyCastleProvider.PROVIDER_NAME) == null) {
            Security.addProvider(new BouncyCastleProvider());
        }
    }

    public static byte[] creteKey() {
        try {
            KeyGenerator keyGenerator = KeyGenerator.getInstance(SM4, BouncyCastleProvider.PROVIDER_NAME);
            keyGenerator.init(128);
            return keyGenerator.generateKey().getEncoded();
        } catch (NoSuchAlgorithmException | NoSuchProviderException e) {
            e.printStackTrace();
        }
        return null;
    }

    /**
     * 加密
     *
     * @param key       对称密钥
     * @param plainData 明文字节
     * @return
     */
    public static byte[] encrypt(byte[] key, byte[] plainData) {
        Cipher cipher = buildCipher(SM4, MODE, PADDING, Cipher.ENCRYPT_MODE, key,
                                    NEED_IV ? ByteUtils.fromHexString(IV_HEX_STRING) : null);
        try {
            byte[] cipherData = cipher.doFinal(plainData);
            return cipherData;
        } catch (IllegalBlockSizeException | BadPaddingException e) {
            e.printStackTrace();
        }
        return null;
    }

    /**
     * 解密
     *
     * @param key        对称密钥
     * @param cipherData 密文字节
     * @return
     */
    public static byte[] decrypt(byte[] key, byte[] cipherData) {
        Cipher cipher = buildCipher(SM4, MODE, PADDING, Cipher.DECRYPT_MODE, key,
                                    NEED_IV ? ByteUtils.fromHexString(IV_HEX_STRING) : null);
        try {
            byte[] plainData = cipher.doFinal(cipherData);
            return plainData;
        } catch (IllegalBlockSizeException | BadPaddingException e) {
            e.printStackTrace();
        }
        return null;
    }

    private static Cipher buildCipher(String algorithmName, String algorithmMode, String algorithmPadding, int cipherMode,
                                      byte[] key, byte[] iv) {
        String algorithmType = String.format("%s/%s/%s", algorithmName, algorithmMode, algorithmPadding);
        SecretKeySpec secretKeySpec = new SecretKeySpec(key, algorithmName);
        try {
            Cipher cipher = Cipher.getInstance(algorithmType, BouncyCastleProvider.PROVIDER_NAME);
            if (iv != null && iv.length != 0) {
                IvParameterSpec ivParameterSpec = new IvParameterSpec(iv);
                cipher.init(cipherMode, secretKeySpec, ivParameterSpec);
            } else {
                cipher.init(cipherMode, secretKeySpec);
            }
            return cipher;
        } catch (NoSuchAlgorithmException | NoSuchProviderException | NoSuchPaddingException | InvalidAlgorithmParameterException | InvalidKeyException e) {
            e.printStackTrace();
        }
        return null;
    }

}
